package njust.csc.aecore.utils

import io.github.oshai.kotlinlogging.KotlinLogging
import org.opencv.core.Core
import org.opencv.core.Mat
import org.opencv.core.MatOfFloat
import org.opencv.imgcodecs.Imgcodecs
import org.opencv.imgproc.Imgproc
import java.io.FileNotFoundException
import java.nio.file.Path
import java.util.*
import kotlin.io.path.exists
import kotlin.io.path.readBytes

private val logger = KotlinLogging.logger {}

/**
 * @Author ZhouHang
 * @Date 24/7/3 周三
 */
object ImageUtils {
    init {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME)
    }

    private fun readImage(imagePath: Path): Mat {
        if (!imagePath.exists()) {
            throw FileNotFoundException("$imagePath 图片不存在")
        }
        return Imgcodecs.imread(imagePath.toAbsolutePath().toString())
    }

    /**
     * @return 越接近 1 说明该图片越空白
     * */
    fun imageFilter(imagePath: Path): Double {
        val image = readImage(imagePath)
        if (image.empty()) {
            logger.info { "$imagePath 图片读取失败" }
            return .0
        }
        val gray = toGray(image)
        image.release()
        return imageFilter(gray)
    }

    private fun imageFilter(imageGray: Mat): Double {
        val size = imageGray.width() * imageGray.height()
        // 二值化
        val imageBinary = Mat()
        Imgproc.threshold(imageGray, imageBinary, 200.0, 255.0, Imgproc.THRESH_BINARY)
        imageGray.release()
        val blankArea = Core.countNonZero(imageBinary)
        imageBinary.release()
        return blankArea / size.toDouble()
    }

    fun calculateSimilarity(image1Path: Path, image2Path: Path): Double {
        val image1 = readImage(image1Path)
        if (image1.empty()) {
            logger.info { "$image1Path 图片读取失败" }
            return .0
        }
        val image2 = readImage(image2Path)
        if (image2.empty()) {
            logger.info { "$image2Path 图片读取失败" }
            return .0
        }
        var image1Gray = toGray(image1)
        var image2Gray = toGray(image2)
        image1.release()
        image2.release()
        // 统一 image1Gray 和 image2Gray 的尺寸
        if (image1Gray.width() > image1Gray.height()) image1Gray = rotate90(image1Gray)
        if (image2Gray.width() > image2Gray.height()) image2Gray = rotate90(image2Gray)
        // 计算相似度
        val similarity = calculateSimilarity(image1Gray, image2Gray)
        image1Gray.release()
        image2Gray.release()
        return similarity
    }

    /**
     * 标准化互相关匹配（TM_CCOEFF_NORMED）计算相似度
     * */
    internal fun calculateSimilarity(image1Gray: Mat, image2Gray: Mat): Double {
        val ssimMat = MatOfFloat()
        Imgproc.matchTemplate(image1Gray, image2Gray, ssimMat, Imgproc.TM_CCOEFF_NORMED)
        val ssimScalar = Core.mean(ssimMat)
        ssimMat.release()
        return ssimScalar.`val`[0]
    }

    /**
     * 转为灰度图像
     * @param image
     * */
    private fun toGray(image: Mat): Mat {
        val imageGray = Mat()
        Imgproc.cvtColor(image, imageGray, Imgproc.COLOR_BGR2GRAY)
        return imageGray
    }

    /**
     * 将给定图片旋转 90 度输出
     * */
    internal fun rotate90(image: Mat): Mat {
        // fastbot 可能会生成一张旋转后的图片，需要转回去
        val transpose = Mat()
        Core.transpose(image, transpose)
        val flip = Mat()
        Core.flip(transpose, flip, 0)
        image.release()
        transpose.release()
        return flip
    }

    /**
     * 对给定图片路径计算 base64 编码
     * @param imagePath Path 需要编码的图片路径
     * @return String 该图片的 base64 编码字符串
     */
    fun base64Encode(imagePath: Path): String {
        return Base64.getEncoder().encodeToString(imagePath.readBytes())
    }

    /**
     * 对给定图片的 base64 字符串解码
     * @param imageString String 需要解码的 base64 字符串
     * @return ByteArray
     */
    fun base64Decode(imageString: String): ByteArray {
        return Base64.getDecoder().decode(imageString)
    }
}
